package chen.yi.component.log.config;

import chen.yi.component.log.utils.ThreadMdcUtil;
import com.alibaba.ttl.spi.TtlEnhanced;
import com.alibaba.ttl.threadpool.TtlExecutors;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.config.BeanPostProcessor;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
import org.springframework.stereotype.Component;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.atomic.AtomicInteger;

@Slf4j
@Component
public class TtlThreadPoolBeanPostProcessor implements BeanPostProcessor {
    private final AtomicInteger wrapCounter = new AtomicInteger();

    @Override
    public Object postProcessAfterInitialization(Object bean, String beanName) {
        try {
            if (bean instanceof ThreadPoolTaskExecutor executor) {
                return handleThreadPoolTaskExecutor(executor, beanName);
            } else if (bean instanceof ExecutorService executorService && !isTtlWrapped(executorService)) {
                return handleExecutorService(executorService, beanName);
            }
        } catch (Exception e) {
            log.error("TTL包装异常: 【{}】", beanName, e);
        }
        return bean;
    }

    private Object handleThreadPoolTaskExecutor(ThreadPoolTaskExecutor executor, String beanName) {
        log.info("包装ThreadPoolTaskExecutor: 【{}】", beanName);
        executor.setTaskDecorator(runnable -> {
            if (runnable instanceof TtlEnhanced) {
                return runnable;
            }
            return ThreadMdcUtil.wrap(runnable);
        });
        wrapCounter.incrementAndGet();
        return executor;
    }

    private Object handleExecutorService(ExecutorService executorService, String beanName) {
        log.info("包装ExecutorService: 【{}】", beanName);
        ExecutorService wrapped = TtlExecutors.getTtlExecutorService(executorService);
        wrapCounter.incrementAndGet();
        return wrapped;
    }

    private boolean isTtlWrapped(ExecutorService executorService) {
        return executorService instanceof TtlEnhanced;
    }
}
